package kpython.backend.ir.canon;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;

import kpython.backend.ir.tree.Assignment;
import kpython.backend.ir.tree.Atom;
import kpython.backend.ir.tree.AtomList;
import kpython.backend.ir.tree.CallBlock;
import kpython.backend.ir.tree.Comparison;
import kpython.backend.ir.tree.Expr;
import kpython.backend.ir.tree.FORBlock;
import kpython.backend.ir.tree.FunctionBlock;
import kpython.backend.ir.tree.Goto;
import kpython.backend.ir.tree.IFBlock;
import kpython.backend.ir.tree.IFStmt;
import kpython.backend.ir.tree.IRNode;
import kpython.backend.ir.tree.IRTreeVisitor;
import kpython.backend.ir.tree.Label;
import kpython.backend.ir.tree.ListElement;
import kpython.backend.ir.tree.Not;
import kpython.backend.ir.tree.Operation;
import kpython.backend.ir.tree.PrintBlock;
import kpython.backend.ir.tree.Relational;
import kpython.backend.ir.tree.Return;
import kpython.backend.ir.tree.SimpleBlock;
import kpython.backend.ir.tree.Stmt;
import kpython.backend.ir.tree.WHILEBlock;

public class Optimizer implements IRTreeVisitor {

	public HashMap<String, Stmt> functions = new HashMap<String, Stmt>();
	public List<String> usedFunctions = new ArrayList<String>();
	
	public SimpleBlock initOptimizer(SimpleBlock root) {
		setFunctions();
		SimpleBlock newRoot = new SimpleBlock();
		for(IRNode node: root.getContent())
		{
			if(node instanceof FunctionBlock) {
				functions.put(((FunctionBlock)node).label.label, (FunctionBlock)node);
				continue;
			}
			newRoot.addContent((IRNode)node.accept(this, null));
		}		
		
		do {
			List<Stmt> functionsToAdd = new LinkedList<Stmt>();
			List<String> usedFunctions2 = new LinkedList<String>();

			for(String function: usedFunctions) {
				functionsToAdd.add(functions.get(function));
				usedFunctions2.add(function);
			}

			for(Stmt stmt: functionsToAdd)
				newRoot.addContent((IRNode)stmt.accept(this, null));

			usedFunctions.removeAll(usedFunctions2);				
		}while(!usedFunctions.isEmpty());
		
		//newRoot.addContent(new WriterFunction());
		//newRoot.addContent(new AccessFunction());
		return newRoot;
	}
	
	private void setFunctions() {
		functions.put("EqEqualFunction", new EqEqualFunction());
		functions.put("EqGreaterFunction", new EqGreaterFunction());
		functions.put("EqLessFunction", new EqLessFunction());
		functions.put("GreaterFunction", new GreaterFunction());
		functions.put("LessFunction", new LessFunction());
		functions.put("DifFunction", new DifFunction());
		functions.put("NotFunction", new NotFunction());
		functions.put("OrFunction", new OrFunction());
		functions.put("AndFunction", new AndFunction());		
		functions.put("DivideFunction", new DivideFunction());
		functions.put("MultiplyFunction", new MultiplyFunction());
		functions.put("PlusFunction", new PlusFunction());
		functions.put("MinusFunction", new MinusFunction());
		functions.put("RestFunction", new RestFunction());
		functions.put("PrintFunction", new PrintFunction());		
	}

	@Override
	public Assignment visit(Assignment node, Object data) {
		node.expr = (Expr)node.expr.accept(this, data);
		node.atom = (Atom)node.atom.accept(this, data);
		return node;
	}

	@Override
	public Atom visit(Atom node, Object data) {				
		return node;
	}

	@Override
	public AtomList visit(AtomList node, Object data) {
		AtomList newList = new AtomList();
		for(IRNode expr: node.getItems())
			newList.addItem((IRNode)expr.accept(this, data));
		return newList;
	}

	@Override
	public CallBlock visit(CallBlock node, Object data) {
		String temp = node.gotoStmt.label;
		if(!usedFunctions.contains(temp))
			usedFunctions.add(temp);	
		return node;
	}

	@Override
	public CallBlock visit(Comparison node, Object data) {
		CallBlock block = new CallBlock();
		String temp = "";
		switch(node.type) {
			case 0:		
				temp = "EqEqualFunction";						
				break;
			case 1:
				temp = "LessFunction";
				break;
			case 2:
				temp = "GreaterFunction";
				break;
			case 3:
				temp = "EqLessFunction";
				break;				
			case 4:
				temp = "EqGreaterFunction";
				break;
			case 5:
				temp = "DifFunction";
				break;
			default:;
		}		
		block.gotoStmt = new Goto();
		block.gotoStmt.label = temp.toLowerCase().replace("function", "");
		block.params.add((Expr)node.left.accept(this, data));
		block.params.add((Expr)node.right.accept(this, data));
		if(!usedFunctions.contains(temp))
			usedFunctions.add(temp);	
		return block;
	}

	@Override
	public Object visit(FORBlock node, Object data) {
		return node;
	}

	@Override
	public FunctionBlock visit(FunctionBlock node, Object data) {				
		node.block = (SimpleBlock)node.block.accept(this, data);
		return node;
	}

	@Override
	public Goto visit(Goto node, Object data) {	
		return node;
	}

	@Override
	public IFBlock visit(IFBlock node, Object data) {	
		if(node.elseBlock != null)
			node.elseBlock = (SimpleBlock)node.elseBlock.accept(this, data);
		List<IFStmt> newStmts = new LinkedList<IFStmt>();
		for(IFStmt stmt : node.stmts)
			newStmts.add((IFStmt)stmt.accept(this, data));
		node.stmts = newStmts;		
		return node;
	}

	@Override
	public IFStmt visit(IFStmt node, Object data) {
		node.expr = (Expr)node.expr.accept(this, data);
		node.block = (SimpleBlock)node.block.accept(this, data);		
		return node;
	}

	@Override
	public Label visit(Label node, Object data) {
		return node;
	}

	@Override
	public CallBlock visit(Operation node, Object data) {
		CallBlock block = new CallBlock();
		String temp = "";
		switch(node.type) {
			case 0:		
				temp = "PlusFunction";						
				break;
			case 1:
				temp = "MinusFunction";
				break;
			case 2:
				temp = "MultiplyFunction";
				break;
			case 3:
				temp = "DivideFunction";
				break;				
			case 4:
				temp = "RestFunction";
				break;			
			default:;
		}		
		block.gotoStmt = new Goto();
		block.gotoStmt.label = temp.toLowerCase().replace("function", "");
		block.params.add((Expr)node.left.accept(this, data));
		block.params.add((Expr)node.right.accept(this, data));
		if(!usedFunctions.contains(temp))
			usedFunctions.add(temp);	
		return block;
	}

	@Override
	public CallBlock visit(PrintBlock node, Object data) {		
		CallBlock block = new CallBlock();
		String temp = "PrintFunction";		
		block.gotoStmt = new Goto();
		block.gotoStmt.label = temp.toLowerCase().replace("function", "");
		block.addExpr((Expr)node.expr.accept(this, data));		
		if(!usedFunctions.contains(temp))
			usedFunctions.add(temp);
		return block;
	}

	@Override
	public CallBlock visit(Relational node, Object data) {
		CallBlock block = new CallBlock();
		String temp = "";
		switch(node.type) {
			case 0:		
				temp = "AndFunction";						
				break;
			case 1:
				temp = "OrFunction";
				break;		
			default:;
		}		
		block.gotoStmt = new Goto();
		block.gotoStmt.label = temp.toLowerCase().replace("function", "");
		block.params.add((Expr)node.left.accept(this, data));
		block.params.add((Expr)node.right.accept(this, data));
		if(!usedFunctions.contains(temp))
			usedFunctions.add(temp);	
		return block;		
	}

	@Override
	public Return visit(Return node, Object data) {
		Return newNode = new Return();
		newNode.expr = (Expr)node.expr.accept(this, data);		
		return newNode;
	}

	@Override
	public SimpleBlock visit(SimpleBlock node, Object data) {
		SimpleBlock newBlock = new SimpleBlock();		
		for(IRNode n: node.getContent())
			newBlock.addContent((IRNode)n.accept(this, data));		
		return newBlock;
	}

	@Override
	public Object visit(WHILEBlock node, Object data) {		
		return node;
	}
	

	@Override
	public CallBlock visit(Not node, Object data) {
		CallBlock block = new CallBlock();		
		String temp = "NotFunction";	
		block.gotoStmt = new Goto();
		block.gotoStmt.label = temp.toLowerCase().replace("function", "");
		block.params.add((Expr)node.expr.accept(this, data));
		if(!usedFunctions.contains(temp))
			usedFunctions.add(temp);	
		return block;
	}
	
	@Override
	public Object visit(ListElement node, Object data) {		
		return node;
	}

	@Override
	public AccessFunction visit(AccessFunction node, Object data) {		
		return node;
	}

	@Override
	public DifFunction visit(DifFunction node, Object data) {
		return node;
	}

	@Override
	public DivideFunction visit(DivideFunction node, Object data) {
		return node;
	}

	@Override
	public EqEqualFunction visit(EqEqualFunction node, Object data) {
		return node;
	}

	@Override
	public EqGreaterFunction visit(EqGreaterFunction node, Object data) {		
		return node;
	}

	@Override
	public EqLessFunction visit(EqLessFunction node, Object data) {
		return node;
	}

	@Override
	public GreaterFunction visit(GreaterFunction node, Object data) {
		return node;
	}

	@Override
	public LessFunction visit(LessFunction node, Object data) {
		return node;
	}

	@Override
	public MinusFunction visit(MinusFunction node, Object data) {
		return node;
	}

	@Override
	public MultiplyFunction visit(MultiplyFunction node, Object data) {
		return node;
	}

	@Override
	public NotFunction visit(NotFunction node, Object data) {
		return node;
	}

	@Override
	public OrFunction visit(OrFunction node, Object data) {
		return node;
	}

	@Override
	public PlusFunction visit(PlusFunction node, Object data) {
		return node;
	}

	@Override
	public PrintFunction visit(PrintFunction node, Object data) {
		return node;
	}

	@Override
	public RestFunction visit(RestFunction node, Object data) {
		return node;
	}

	@Override
	public WriterFunction visit(WriterFunction node, Object data) {
		return node;
	}

	@Override
	public AndFunction visit(AndFunction node, Object data) {
		return node;
	}
}
